﻿#include "ursaencrypt.h"

#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>

#include <QString>
#include <QDebug>

URSAEncrypt::URSAEncrypt()
{
}

URSAEncrypt::~URSAEncrypt()
{
}

void URSAEncrypt::generateRSAKeys(int length, const QString &publicKeyFile, const QString &privateKeyFile)
{
    RSA* rsa = RSA_new();
    BIGNUM* bne = BN_new();
    unsigned long public_exponent = RSA_F4;

    if (BN_set_word(bne, public_exponent) != 1)
    {
        qCritical() << "Error setting public exponent.";
        return;
    }

    if (RSA_generate_key_ex(rsa, length, bne, nullptr) != 1)
    {
        qCritical() << "Error generating RSA key.";
        return;
    }

    BIO* private_out = BIO_new_file(privateKeyFile.toStdString().c_str(), "w");
    PEM_write_bio_RSAPrivateKey(private_out, rsa, nullptr, nullptr, 0, nullptr, nullptr);
    BIO_free(private_out);

    BIO* public_out = BIO_new_file(publicKeyFile.toStdString().c_str(), "w");
    PEM_write_bio_RSA_PUBKEY(public_out, rsa);
    BIO_free(public_out);

    RSA_free(rsa);
    BN_free(bne);
}

QByteArray URSAEncrypt::encryptByPubkeyString(const QByteArray &message, const QByteArray &pubkey)
{
    BIO* in = BIO_new_mem_buf((void*)pubkey.data(), -1);
    if (in == nullptr)
    {
        qCritical() << "BIO_new_mem_buf failed";
        return QByteArray();
    }

    RSA* rsa = PEM_read_bio_RSA_PUBKEY(in, nullptr, nullptr, nullptr);
    BIO_free(in);
    if (rsa == nullptr)
    {
        qCritical() << "PEM_read_bio_RSA_PUBKEY failed";
        return QByteArray();
    }

    int size = RSA_size(rsa);
    QByteArray encrypt_data;
    encrypt_data.resize(size);
    int ret = RSA_public_encrypt(message.size(), (unsigned char*)message.data(), (unsigned char*)encrypt_data.data(), rsa, RSA_PKCS1_PADDING);
    RSA_free(rsa);
    if (ret == -1)
    {
        qCritical() << "RSA_public_encrypt failed";
        return QByteArray();
    }

    return encrypt_data;
}

QByteArray URSAEncrypt::encryptByPubkeyFile(const QByteArray &message, const std::string &pub_filename)
{
    BIO* in = BIO_new(BIO_s_file());
    if (in == nullptr)
    {
        qCritical() << "BIO_new failed";
        return QByteArray();
    }
    BIO_read_filename(in, pub_filename.c_str());

    RSA* rsa = PEM_read_bio_RSA_PUBKEY(in, nullptr, nullptr, nullptr);
    BIO_free(in);
    if (rsa == nullptr)
    {
        qCritical() << "PEM_read_bio_RSA_PUBKEY failed";
        return QByteArray();
    }
    int size = RSA_size(rsa);
    QByteArray encrypt_data;
    encrypt_data.resize(size);
    int ret = RSA_public_encrypt(message.size(), (unsigned char*)message.data(), (unsigned char*)encrypt_data.data(), rsa, RSA_PKCS1_PADDING);
    RSA_free(rsa);
    if (ret == -1)
    {
        qCritical() << "RSA_public_encrypt failed";
        return QByteArray();
    }
    return encrypt_data;
}

QByteArray URSAEncrypt::encryptByPrikeyString(const QByteArray &message, const QByteArray &prikey)
{
    BIO* in = BIO_new_mem_buf((void*)prikey.data(), -1);
    if (in == nullptr)
    {
        qCritical() << "BIO_new_mem_buf failed";
        return QByteArray();
    }

    RSA* rsa = PEM_read_bio_RSAPrivateKey(in, nullptr, nullptr, nullptr);
    BIO_free(in);
    if (rsa == nullptr)
    {
        qCritical() << "PEM_read_bio_RSAPrivateKey failed";
        return QByteArray();
    }

    int size = RSA_size(rsa);
    QByteArray encrypt_data;
    encrypt_data.resize(size);
    int ret = RSA_private_encrypt(message.size(), (unsigned char*)message.data(), (unsigned char*)encrypt_data.data(), rsa, RSA_PKCS1_PADDING);
    RSA_free(rsa);
    if (ret == -1)
    {
        qCritical() << "RSA_private_encrypt failed";
        return QByteArray();
    }

    return encrypt_data;
}

QByteArray URSAEncrypt::encryptByPrikeyFile(const QByteArray &message, const std::string &pri_file)
{
    BIO* in = BIO_new(BIO_s_file());
    if (in == nullptr)
    {
        qCritical() << "BIO_new failed";
        return QByteArray();
    }
    BIO_read_filename(in, pri_file.c_str());

    RSA* rsa = PEM_read_bio_RSAPrivateKey(in, nullptr, nullptr, nullptr);
    BIO_free(in);
    if (rsa == nullptr)
    {
        qCritical() << "PEM_read_bio_RSAPrivateKey failed";
        return QByteArray();
    }
    int size = RSA_size(rsa);
    QByteArray encrypt_data;
    encrypt_data.resize(size);
    int ret = RSA_private_encrypt(message.size(), (unsigned char*)message.data(), (unsigned char*)encrypt_data.data(), rsa, RSA_PKCS1_PADDING);
    RSA_free(rsa);
    if (ret == -1)
    {
        qCritical() << "RSA_private_encrypt failed";
        return QByteArray();
    }
    return encrypt_data;
}

QByteArray URSAEncrypt::decryptByPubkeyString(const QByteArray &cipher, const QByteArray &pubkey)
{
    BIO* in = BIO_new_mem_buf((void*)pubkey.data(), -1);
    if (in == nullptr)
    {
        qCritical() << "BIO_new_mem_buf failed";
        return "";
    }

    RSA* rsa = PEM_read_bio_RSA_PUBKEY(in, nullptr, nullptr, nullptr);
    BIO_free(in);
    if (rsa == nullptr)
    {
        qCritical() << "PEM_read_bio_RSA_PUBKEY failed";
        return "";
    }

    int size = RSA_size(rsa);
    QByteArray data;
    data.resize(size);
    int ret = RSA_public_decrypt(cipher.size(), (unsigned char*)cipher.data(), (unsigned char*)data.data(), rsa, RSA_PKCS1_PADDING);
    RSA_free(rsa);
    if (ret == -1)
    {
        qCritical() << "RSA_public_decrypt failed";
        return "";
    }

    return data.left(ret);
}

QByteArray URSAEncrypt::decryptByPubkeyFile(const QByteArray &cipher, const std::string &pub_filename)
{
    BIO* in = BIO_new(BIO_s_file());
    if (in == nullptr)
    {
        qCritical() << "BIO_new failed";
        return "";
    }
    BIO_read_filename(in, pub_filename.c_str());

    RSA* rsa = PEM_read_bio_RSA_PUBKEY(in, nullptr, nullptr, nullptr);
    BIO_free(in);
    if (rsa == nullptr)
    {
        qCritical() << "PEM_read_bio_RSA_PUBKEY failed";
        return "";
    }

    int size = RSA_size(rsa);
    QByteArray data;
    data.resize(size);
    int ret = RSA_public_decrypt(cipher.size(), (unsigned char*)cipher.data(), (unsigned char*)data.data(), rsa, RSA_PKCS1_PADDING);
    RSA_free(rsa);
    if (ret == -1)
    {
        qCritical() << "RSA_public_decrypt failed";
        return "";
    }

    return data.left(ret);
}

QByteArray URSAEncrypt::decryptByPrikeyString(const QByteArray &cipher, const QByteArray &prikey)
{
    BIO* in = BIO_new_mem_buf((void*)prikey.data(), -1);
    if (in == nullptr)
    {
        qCritical() << "BIO_new_mem_buf failed";
        return "";
    }

    RSA* rsa = PEM_read_bio_RSAPrivateKey(in, nullptr, nullptr, nullptr);
    BIO_free(in);
    if (rsa == nullptr)
    {
        qCritical() << "PEM_read_bio_RSAPrivateKey failed";
        return "";
    }

    int size = RSA_size(rsa);
    QByteArray data;
    data.resize(size);
    int ret = RSA_private_decrypt(cipher.size(), (unsigned char*)cipher.data(), (unsigned char*)data.data(), rsa, RSA_PKCS1_PADDING);
    RSA_free(rsa);
    if (ret == -1)
    {
        qCritical() << "RSA_private_decrypt failed";
        return "";
    }

    return data.left(ret);
}

QByteArray URSAEncrypt::decryptByPrikeyFile(const QByteArray &cipher, const std::string &pri_file)
{
    BIO* in = BIO_new(BIO_s_file());
    if (in == NULL)
    {
        qCritical() << "BIO_new failed";
        return QByteArray();
    }
    BIO_read_filename(in, pri_file.c_str());

    RSA* rsa = PEM_read_bio_RSAPrivateKey(in, nullptr, nullptr, nullptr);
    BIO_free(in);
    if (rsa == nullptr)
    {
        qCritical() << "PEM_read_bio_RSAPrivateKey failed";
        return QByteArray();
    }

    int size = RSA_size(rsa);
    QByteArray data;
    data.resize(size);
    int ret = RSA_private_decrypt(cipher.size(), (unsigned char*)cipher.data(), (unsigned char*)data.data(), rsa, RSA_PKCS1_PADDING);
    RSA_free(rsa);
    if (ret == -1)
    {
        qCritical() << "RSA_private_decrypt failed";
        return QByteArray();
    }

    return data.left(ret);
}

QString URSAEncrypt::GetLastError()
{
    unsigned long lastError = ERR_peek_last_error();
    if (lastError == 0)
    {
        return "";
    }
    char errorString[256];
    ERR_error_string_n(lastError, errorString, sizeof(errorString));
    return errorString;
}
